# import tkinter as tk
import numpy as np
import random
from joblib import Parallel, delayed
import json
import subprocess
import os
import glob
# import tkinter as tk
# from tkinter import simpledialog

def call_func(input_func, populations, kwargs):
    return input_func(
        populations,
        **kwargs  # 解包 extra_param 作為關鍵字參數
    )

def get_latest_txt_file(directory):
    # 获取目录中所有 .txt 文件的完整路径
    txt_files = glob.glob(os.path.join(directory, 'benchmark_result*.txt'))
    
    # 如果没有找到 .txt 文件，返回 None
    if not txt_files:
        return None
    
    # 找到最新的 .txt 文件
    latest_file = max(txt_files, key=os.path.getctime)
    return latest_file

# 假設 call_func 和 pbest 是已經定義好的
def parallel_call_func(func, p, trans_params):
    return call_func(func, p, trans_params)

# 定義並行計算函數
def compute_populations_fitness(pbest_combinations, func, trans_params):
    output = Parallel(n_jobs=-1)(delayed(parallel_call_func)(func, combo, trans_params) for combo in pbest_combinations)
    print(f'parallel operation output:', output)
    output_array = np.array(output)
    print(output_array.shape)
    throughput, decode_latency, prefill_latency = output_array[:, 0], output_array[:, 1], output_array[:, 2]

    return np.array(throughput), np.array(decode_latency), np.array(prefill_latency)

def compute_minserver(combinations, model_name, concurrency, request_rate):

    for i, combo in enumerate(combinations):
        with open(f'../run_test_sh/conf_case/{model_name}_config.json', 'r') as file:
            data = json.load(file)
        data["BackendConfig"]["ScheduleConfig"]["maxBatchSize"] = int(combo['Decode BatchSize'])
        data["BackendConfig"]["ScheduleConfig"]["maxPrefillBatchSize"] = int(combo['Prefill BatchSize'])
        data["BackendConfig"]["ScheduleConfig"]["prefillTimeMsPerReq"] = int(combo['SelectBatch Prefill Delay Tolerance'])
        with open(f'../run_test_sh/conf_case/{model_name}/config_{i}.json', 'w') as file:
            json.dump(data, file)

    with subprocess.Popen(['bash', '../run_test_sh/multi_run_multi_env.sh', model_name, str(concurrency), str(request_rate)], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True) as process:
        for line in process.stdout:
            print(line.rstrip())
        # 读取错误输出（如果有）
        for line in process.stderr:
            print(f"Error: {line}")
    
    benchmark_result = get_latest_txt_file("./benchmark/")


    get_result = {"FirstTokenTime":[], "DecodeTime" : [], "GenerateSpeed": []}
    key_words = ['FirstTokenTime','DecodeTime','GenerateSpeed']
    with open(benchmark_result, 'r') as file:
        for line_number, line in enumerate(file, start=1):
            # 按 '|' 分割行内容并去除每个部分的首尾空格
            parts = [part.strip() for part in line.strip().split('|')]
            if len(parts) >1 and parts[1] in key_words:
                get_result[parts[1]].append(float(parts[2].split(' ')[0]))
    print(get_result)
    return  np.array(get_result['GenerateSpeed']),  np.array(get_result['DecodeTime']),  np.array(get_result['FirstTokenTime'])

def generate_populations_2bit(bounds, num_populations, gene_length):
    populations_2bit = {}

    values = list(bounds.values())  # 获取字典的值列表
    keys = list(bounds.keys())      # 获取字典的键列表

    for i in range(len(bounds)):
        populations_2bit[keys[i]] = np.random.randint(2, size=(num_populations, gene_length))
    return populations_2bit

def decode(populations, lower_bound, upper_bound, gene_length):
    return populations.dot(2 ** np.arange(gene_length)[::-1]) *(upper_bound-lower_bound) / float(2**gene_length-1) + lower_bound

def decode_populations_2bit(bounds, populations):
    '''
    {'Prefill BatchSize': array([[0, 1, 0, 1, 0],
                                [0, 0, 0, 0, 1],
                                [1, 0, 0, 1, 1],
                                [0, 0, 1, 0, 0]]),
     'Decode BatchSize': array([[0, 1, 1, 1, 0],
                                [0, 1, 0, 1, 1],
                                [0, 1, 0, 1, 0],
                                [0, 1, 1, 1, 0]])
    }
    '''

    populations_decoded = {}

    values = list(bounds.values())  # 获取字典的值列表
    keys = list(bounds.keys())      # 获取字典的键列表

    for i in range(len(bounds)):
        lower_bound, upper_bound = values[i]  # 直接通过索引访问边界值
        populations_decoded[keys[i]] = [decode(item, lower_bound, upper_bound, len(item)) for item in populations[keys[i]]]
    return populations_decoded

def round_populations(populations):
    # 创建一个新的字典来存储取整后的粒子值
    rounded_populations = {}

    for key, value in populations.items():
        # 确保 value 是一个 NumPy 数组
        value_array = np.array(value)

        # 对 NumPy 数组中的值进行取整
        rounded_value = np.floor(value_array).astype(int)

        # 将处理后的值存回新字典中
        rounded_populations[key] = rounded_value

    return rounded_populations

def penalty_functions(constrain_params, hard_constrain):
    # 对每个 hard_constrain 计算惩罚
    fitness_penalty = []
    for i in range(len(constrain_params)):
        penalty = (max(0, (constrain_params[i] - hard_constrain) / hard_constrain)) ** 2
        fitness_penalty.append(penalty)  # 累加惩罚
    return fitness_penalty

def fitness_functions(criteria, constrain_params, hard_constrain):
    # 计算初始 fitness
    fitness = 1 / criteria

    # 遍历 constrain_params 和其对应的 hard_constrain
    print(constrain_params)
    print(fitness)
    for idx, constrain_param in enumerate(constrain_params):
        # 针对每个约束，计算惩罚并应用到 fitness 上
        print("constrain_param: ",constrain_param)
        print("hard_constrain: ",hard_constrain[idx])
        fitness += penalty_functions(constrain_param, hard_constrain[idx])
        print("fitness: ",fitness)
    return fitness

def penalty_function(criteria, hard_constrain):
    fitness_penalty = (max(0, (criteria - hard_constrain)/hard_constrain)) ** 2
    return fitness_penalty

def fitness_function(criteria, constrain_params, hard_constrain):
    '''
    :param criteria: [3458.52059309 3454.51253015 3431.46376627 3319.62259369]
    :param constrain_params: [0.06994235 0.07095963 0.06717097 0.05938756]
    :param hard_constrain: 0.05
    :return:
    '''
    # 计算初始 fitness
    fitness = 1 / criteria

    # 遍历 constrain_params 和其索引
    for idx, constrain_param in enumerate(constrain_params):
        fitness[idx] += penalty_function(constrain_param, hard_constrain)

    return fitness

# 选择过程
def selection(fitness_list, num_parents):
    # 根据适应度排序，选择适应度最低的个体作为父母
    # num_parents: 选择的父母数量
    sorted_indices = sorted(range(len(fitness_list)), key=lambda x: fitness_list[x])
    return sorted_indices[:num_parents]

# 交叉操作
def crossover(populations, crossover_rate=0.8):
    offspring = {}
    keys = list(populations.keys())

    for key in keys:
        population = populations[key]
        offspring[key] = np.copy(population)  # 复制以保留父代

        for i in range(0, len(population), 2):
            if i + 1 < len(population) and np.random.rand() < crossover_rate:
                crossover_point = np.random.randint(1, population.shape[1])  # 随机选择交叉点
                offspring[key][i, crossover_point:], offspring[key][i + 1, crossover_point:] = \
                    population[i + 1, crossover_point:], population[i, crossover_point:]

    return offspring

# 变异操作
def mutate(populations, mutation_rate=0.2):
    for key, population in populations.items():
        for i in range(len(population)):
            for j in range(population.shape[1]):
                if np.random.rand() < mutation_rate:
                    population[i, j] = 1 - population[i, j]  # 翻转位

    return populations

def merge_populations(parents, offspring):
    # 创建一个新的字典来存储合并后的种群
    merged_population = {}

    # 对于每个键值，合并父代和子代
    for key in parents.keys():
        merged_population[key] = np.vstack((parents[key], offspring[key]))

    return merged_population

def genetic_altorithm(func, bounds, hard_constrains, num_populations, gene_length, num_iterations, model_name, concurrency, request_rate, isprefill_constrained):
    # 初始化种群，使用2进制编码
    populations_2bit = generate_populations_2bit(bounds, num_populations, gene_length)

    # 将2进制解码成10进制，并取整
    populations = decode_populations_2bit(bounds, populations_2bit)
    populations = round_populations(populations)

    # 生成所有可能的组合
    keys = list(populations.keys()) # prefill或者decode
    values_list = [populations[key] for key in keys] # [[...], [...]]
    combinations = [dict(zip(keys, [values_list[i][j] for i in range(len(keys))])) for j in range(len(values_list[0]))]
    print(f'种群初始位置:', combinations)

    # 计算出吞吐量和时延
    # throughput, decode_latency, prefill_latency = compute_populations_fitness(combinations, func, trans_params)
    throughput, decode_latency, prefill_latency = compute_minserver(combinations, model_name, concurrency, request_rate)

    # 计算适应度，在该场景下适应度越小越好
    if isprefill_constrained:
        fitness_list = fitness_functions(throughput, [decode_latency, prefill_latency], hard_constrains)
    else:
        fitness_list = fitness_function(throughput, decode_latency, hard_constrains)
    # fitness_list = fitness_function(throughput, decode_latency, hard_constrains)

    # 找到最小的适应度值的索引
    best_fitness_index = np.argmin(fitness_list)
    gbest = {key: values[best_fitness_index] for key, values in populations.items()}
    gbest_fitness = np.min(fitness_list)
    gbest_output = throughput[best_fitness_index]
    g_decode_latency = decode_latency[best_fitness_index]
    g_prefill_latency = prefill_latency[best_fitness_index]

    with open(f"{model_name}_genetic_result.txt","w") as f:
        f.write(f"初次最优吞吐: {gbest_output}\n")
        f.write(f"初次最优适应性: {gbest_fitness}\n")
        f.write(f"初次最优粒子参数: {gbest}\n")
        f.write(f"初次最优性能对应LPOT: {g_decode_latency}\n")
        f.write(f"初次最优性能对应TTFT: {g_prefill_latency}\n\n")

    print(f'初次最优吞吐:', gbest_output)
    print(f'初次最优适应性:', gbest_fitness)
    print(f'初次最优性能对应LPOT:', g_decode_latency)
    print(f'初次最优性能对应TTFT:', g_prefill_latency)

    # 迭代优化
    for _ in range(num_iterations):
        # 选择
        selected_indices = selection(fitness_list, 4)
        selected_populations_2bit = {key: value[selected_indices] for key, value in populations_2bit.items()}

        # 执行交叉
        crossover_populations_2bit = crossover(selected_populations_2bit)

        # 合并被选择的个体和交叉后的个体
        populations_2bit = merge_populations(selected_populations_2bit, crossover_populations_2bit)

        # 执行变异
        populations_2bit = mutate(populations_2bit)

        # 将2进制解码成10进制，并取整
        populations = decode_populations_2bit(bounds, populations_2bit)
        populations = round_populations(populations)

        # 生成所有可能的组合
        keys = list(populations.keys()) # prefill或者decode
        values_list = [populations[key] for key in keys] # [[...], [...]]
        combinations = [dict(zip(keys, [values_list[i][j] for i in range(len(keys))])) for j in range(len(values_list[0]))]

        print(f'{_ + 1}次种群位置:', combinations)

        # 评估新位置
        # throughput, decode_latency, prefill_latency = compute_populations_fitness(combinations, func, trans_params)
        throughput, decode_latency, prefill_latency = compute_minserver(combinations, model_name, concurrency, request_rate)

        # fitness_list = fitness_function(throughput, decode_latency, hard_constrains)
        if isprefill_constrained:
            fitness_list = fitness_functions(throughput, [decode_latency, prefill_latency], hard_constrains)
        else:
            fitness_list = fitness_function(throughput, decode_latency, hard_constrains)

        updated_fitness_index = np.argmin(fitness_list)

        if np.min(fitness_list) < gbest_fitness:
            gbest = {key: values[updated_fitness_index] for key, values in populations.items()}
            gbest_fitness = np.min(fitness_list)
            gbest_output = throughput[updated_fitness_index]
            g_decode_latency = decode_latency[updated_fitness_index]
            g_prefill_latency = prefill_latency[updated_fitness_index]

        with open(f"{model_name}_genetic_result.txt","a") as f:
            f.write(f"{_ + 1}次全局最优吞吐: {gbest_output}\n")
            f.write(f"{_ + 1}次全局最优适应性: {gbest_fitness}\n")
            f.write(f"{_ + 1}次最优粒子参数: {gbest}\n")
            f.write(f"{_ + 1}次最优性能对应LPOT: {g_decode_latency}\n")
            f.write(f"{_ + 1}次最优性能对应TTFT: {g_prefill_latency}\n")

        print(f'{_ + 1}次全局最优吞吐:', gbest_output)
        print(f'{_ + 1}次全局最优适应性:', gbest_fitness)
        print(f'{_ + 1}次最优粒子参数:', gbest)
        print(f'{_ + 1}次最优性能对应LPOT:', g_decode_latency)
        print(f'{_ + 1}次最优性能对应TTFT:', g_prefill_latency)
 
   return gbest, gbest_output, g_decode_latency, g_prefill_latency

# def get_genetic_input():
#     def submit_data():
#         try:
#             # 从 Entry 组件获取值
#             nonlocal bounds, hard_constrains, num_populations, gene_length
#             bounds = {
#                 'Prefill BatchSize': [
#                     float(pref_batch_lower_entry.get()),
#                     float(pref_batch_upper_entry.get())
#                 ],
#                 'Decode BatchSize': [
#                     float(decode_batch_lower_entry.get()),
#                     float(decode_batch_upper_entry.get())
#                 ],
#                 'SelectBatch Prefill Delay Tolerance': [
#                     float(prefill_delay_tolerance_lower_entry.get()),
#                     float(prefill_delay_tolerance_upper_entry.get())
#                 ]
#             }
#             hard_constrains = float(hard_constrains_entry.get())
#             num_populations = int(num_populations_entry.get())
#             gene_length = int(gene_length_entry.get())

#             # 关闭窗口
#             window.destroy()

#         except ValueError:
#             messagebox.showerror("输入错误", "请输入有效的数字。")

#     # 创建主窗口
#     root = tk.Tk()
#     root.withdraw()  # 隐藏主窗口

#     # 创建自定义窗口
#     window = tk.Toplevel(root)
#     window.title("输入参数")

#     # 默认值
#     default_values = {
#         'pref_batch_lower': '100',
#         'pref_batch_upper': '500',
#         'decode_batch_lower': '100',
#         'decode_batch_upper': '500',
#         'prefill_token_delay_tolerance_lower': '200',
#         'prefill_token_delay_tolerance_upper': '600',
#         'hard_constrains': '0.05',
#         'num_populations': '8',
#         'gene_length': '12'
#     }

#     # 初始化返回值
#     bounds = {}
#     hard_constrains = 0
#     num_populations = 0
#     gene_length = 0

#     # 创建并放置控件
#     tk.Label(window, text="Prefill BatchSize 下界:").grid(row=0, column=0, padx=10, pady=5)
#     pref_batch_lower_entry = tk.Entry(window)
#     pref_batch_lower_entry.insert(0, default_values['pref_batch_lower'])
#     pref_batch_lower_entry.grid(row=0, column=1, padx=10, pady=5)

#     tk.Label(window, text="Prefill BatchSize 上界:").grid(row=1, column=0, padx=10, pady=5)
#     pref_batch_upper_entry = tk.Entry(window)
#     pref_batch_upper_entry.insert(0, default_values['pref_batch_upper'])
#     pref_batch_upper_entry.grid(row=1, column=1, padx=10, pady=5)

#     tk.Label(window, text="Decode BatchSize 下界:").grid(row=2, column=0, padx=10, pady=5)
#     decode_batch_lower_entry = tk.Entry(window)
#     decode_batch_lower_entry.insert(0, default_values['decode_batch_lower'])
#     decode_batch_lower_entry.grid(row=2, column=1, padx=10, pady=5)

#     tk.Label(window, text="Decode BatchSize 上界:").grid(row=3, column=0, padx=10, pady=5)
#     decode_batch_upper_entry = tk.Entry(window)
#     decode_batch_upper_entry.insert(0, default_values['decode_batch_upper'])
#     decode_batch_upper_entry.grid(row=3, column=1, padx=10, pady=5)

#     tk.Label(window, text="Prefill Token Delay Tolerance 下界:").grid(row=4, column=0, padx=10, pady=5)
#     prefill_delay_tolerance_lower_entry = tk.Entry(window)
#     prefill_delay_tolerance_lower_entry.insert(0, default_values['prefill_token_delay_tolerance_lower'])
#     prefill_delay_tolerance_lower_entry.grid(row=4, column=1, padx=10, pady=5)

#     tk.Label(window, text="Prefill Token Delay Tolerance 上界:").grid(row=5, column=0, padx=10, pady=5)
#     prefill_delay_tolerance_upper_entry = tk.Entry(window)
#     prefill_delay_tolerance_upper_entry.insert(0, default_values['prefill_token_delay_tolerance_upper'])
#     prefill_delay_tolerance_upper_entry.grid(row=5, column=1, padx=10, pady=5)

#     tk.Label(window, text="lpot时延边界（秒）:").grid(row=6, column=0, padx=10, pady=5)
#     hard_constrains_entry = tk.Entry(window)
#     hard_constrains_entry.insert(0, default_values['hard_constrains'])
#     hard_constrains_entry.grid(row=6, column=1, padx=10, pady=5)

#     tk.Label(window, text="种群数量:").grid(row=7, column=0, padx=10, pady=5)
#     num_populations_entry = tk.Entry(window)
#     num_populations_entry.insert(0, default_values['num_populations'])
#     num_populations_entry.grid(row=7, column=1, padx=10, pady=5)

#     tk.Label(window, text="基因长度:").grid(row=8, column=0, padx=10, pady=5)
#     gene_length_entry = tk.Entry(window)
#     gene_length_entry.insert(0, default_values['gene_length'])
#     gene_length_entry.grid(row=8, column=1, padx=10, pady=5)

#     # 提交按钮
#     submit_button = tk.Button(window, text="提交", command=submit_data)
#     submit_button.grid(row=9, column=0, columnspan=2, pady=10)

#     # 等待窗口关闭
#     root.wait_window(window)

#     # 返回用户输入的数据
#     return bounds, hard_constrains, num_populations, gene_length

# def get_genetic_input2():
#     def submit_data():
#         try:
#             # 从 Entry 组件获取值
#             nonlocal bounds, hard_constrains, num_populations, gene_length
#             bounds = {
#                 'Prefill BatchSize': [
#                     float(pref_batch_lower_entry.get()),
#                     float(pref_batch_upper_entry.get())
#                 ],
#                 'Decode BatchSize': [
#                     float(decode_batch_lower_entry.get()),
#                     float(decode_batch_upper_entry.get())
#                 ]
#             }
#             hard_constrains = float(hard_constrains_entry.get())
#             num_populations = int(num_populations_entry.get())
#             gene_length = int(gene_length_entry.get())

#             # 关闭窗口
#             window.destroy()

#         except ValueError:
#             messagebox.showerror("输入错误", "请输入有效的数字。")

#     # 创建主窗口
#     root = tk.Tk()
#     root.withdraw()  # 隐藏主窗口

#     # 创建自定义窗口
#     window = tk.Toplevel(root)
#     window.title("输入参数")

#     # 默认值
#     default_values = {
#         'pref_batch_lower': '100',
#         'pref_batch_upper': '500',
#         'decode_batch_lower': '100',
#         'decode_batch_upper': '500',
#         'hard_constrains': '0.05',
#         'num_populations': '8',
#         'gene_length': '12'
#     }

#     # 初始化返回值
#     bounds = {}
#     hard_constrains = 0
#     num_populations = 0
#     gene_length = 0

#     # 创建并放置控件
#     tk.Label(window, text="Prefill BatchSize 下界:").grid(row=0, column=0, padx=10, pady=5)
#     pref_batch_lower_entry = tk.Entry(window)
#     pref_batch_lower_entry.insert(0, default_values['pref_batch_lower'])
#     pref_batch_lower_entry.grid(row=0, column=1, padx=10, pady=5)

#     tk.Label(window, text="Prefill BatchSize 上界:").grid(row=1, column=0, padx=10, pady=5)
#     pref_batch_upper_entry = tk.Entry(window)
#     pref_batch_upper_entry.insert(0, default_values['pref_batch_upper'])
#     pref_batch_upper_entry.grid(row=1, column=1, padx=10, pady=5)

#     tk.Label(window, text="Decode BatchSize 下界:").grid(row=2, column=0, padx=10, pady=5)
#     decode_batch_lower_entry = tk.Entry(window)
#     decode_batch_lower_entry.insert(0, default_values['decode_batch_lower'])
#     decode_batch_lower_entry.grid(row=2, column=1, padx=10, pady=5)

#     tk.Label(window, text="Decode BatchSize 上界:").grid(row=3, column=0, padx=10, pady=5)
#     decode_batch_upper_entry = tk.Entry(window)
#     decode_batch_upper_entry.insert(0, default_values['decode_batch_upper'])
#     decode_batch_upper_entry.grid(row=3, column=1, padx=10, pady=5)

#     tk.Label(window, text="lpot时延边界（秒）:").grid(row=4, column=0, padx=10, pady=5)
#     hard_constrains_entry = tk.Entry(window)
#     hard_constrains_entry.insert(0, default_values['hard_constrains'])
#     hard_constrains_entry.grid(row=4, column=1, padx=10, pady=5)

#     tk.Label(window, text="种群数量:").grid(row=5, column=0, padx=10, pady=5)
#     num_populations_entry = tk.Entry(window)
#     num_populations_entry.insert(0, default_values['num_populations'])
#     num_populations_entry.grid(row=5, column=1, padx=10, pady=5)

#     tk.Label(window, text="基因长度:").grid(row=6, column=0, padx=10, pady=5)
#     gene_length_entry = tk.Entry(window)
#     gene_length_entry.insert(0, default_values['gene_length'])
#     gene_length_entry.grid(row=6, column=1, padx=10, pady=5)

#     # 提交按钮
#     submit_button = tk.Button(window, text="提交", command=submit_data)
#     submit_button.grid(row=9, column=0, columnspan=2, pady=10)

#     # 等待窗口关闭
#     root.wait_window(window)

#     # 返回用户输入的数据
#     return bounds, hard_constrains, num_populations, gene_length

if __name__ == "__main__":
    get_genetic_input2()
#%%

